# coding: utf-8
# File   :  mysql_control
# Author :  wangsheng
# @Time  :  2025/5/8 下午2:26
import pymysql
from dbutils.pooled_db import PooledDB
from typing import Union, List, Dict, Any, Optional
from pymysql.cursors import DictCursor  # 显式导入 DictCursor
from loguru import logger
from utils import config


class MySQLControl:
    """
    MySQL 数据库控制工具类（线程安全）
    功能：增删改查 + 事务支持 + 连接池管理
    依赖：PyMySQL + DBUtils
    """

    def __init__(self,
                 host: str = config.mysql_db.host,
                 port: int = config.mysql_db.port,
                 user: str = config.mysql_db.user,
                 password: str = config.mysql_db.password,
                 database: str = config.mysql_db.database,
                 pool_size: int = 5):
        """
        初始化连接池
        :param pool_size: 连接池大小 (默认5个连接)
        """
        self.pool = PooledDB(
            creator=pymysql,
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            maxconnections=pool_size,
            cursorclass=DictCursor  # 返回字典格式结果
        )

    def _get_connection(self):
        """从连接池获取连接（自动归还）"""
        return self.pool.connection()

    def execute_query(self,
                      sql: str,
                      params: Union[tuple, dict] = None) -> List[Dict[str, Any]]:
        """
        执行查询操作（SELECT）
        :return: 结果列表（每行转为字典）
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    r = cursor.execute(sql, params)
                    logger.debug(f"SQL: {sql}")
                    logger.debug(f"Params: {params}")
                    return cursor.fetchall()
        except pymysql.Error as e:
            logger.error(f"查询执行失败: {e}\nSQL: {sql}\nParams: {params}")
            raise

    def execute_update(self,
                       sql: str,
                       params: Union[tuple, dict] = None) -> int:
        """
        执行更新操作（INSERT/UPDATE/DELETE）
        :return: 影响的行数
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    affected_rows = cursor.execute(sql, params)
                    conn.commit()
                    return affected_rows
        except pymysql.Error as e:
            conn.rollback()
            logger.error(f"更新执行失败: {e}\nSQL: {sql}\nParams: {params}")
            raise

    def insert_one(self,
                   table: str,
                   data: Dict[str, Any]) -> int:
        """
        插入单条数据
        :return: 插入行的自增ID（如果有）
        """
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"

        with self._get_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(sql, tuple(data.values()))
                conn.commit()
                return cursor.lastrowid

    def batch_insert(self,
                     table: str,
                     data_list: List[Dict[str, Any]]) -> int:
        """
        批量插入数据
        :return: 影响的行数
        """
        if not data_list:
            return 0

        columns = ', '.join(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(data_list[0]))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        values = [tuple(item.values()) for item in data_list]

        with self._get_connection() as conn:
            with conn.cursor() as cursor:
                affected_rows = cursor.executemany(sql, values)
                conn.commit()
                return affected_rows

    def execute_transaction(self,
                            sql_actions: List[Dict[str, Union[str, tuple, dict]]]) -> bool:
        """
        执行事务（原子性操作）
        :param sql_actions: 每个元素为 {'sql': str, 'params': tuple/dict}
        :return: 是否成功
        """
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cursor:
                    for action in sql_actions:
                        cursor.execute(action['sql'], action.get('params'))
                    conn.commit()
                    return True
        except pymysql.Error as e:
            conn.rollback()
            logger.error(f"事务执行失败: {e}")
            return False

    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在"""
        sql = """
        SELECT COUNT(*) as count 
        FROM debugrmation_schema.tables 
        WHERE table_schema = DATABASE() AND table_name = %s
        """
        result = self.execute_query(sql, (table_name,))
        return result[0]['count'] > 0

    def close(self):
        """关闭连接池"""
        self.pool.close()
        logger.debug("MySQL连接池已关闭")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# 使用示例
if __name__ == "__main__":

    # 创建工具实例
    mysql = MySQLControl()

    try:

        # 2. 插入单条数据
        # user_id = mysql.insert_one('users', {
        #     'name': 'Alice',
        #     'age': 25,
        #     'email': 'alice@example.com'
        # })
        # print(f"插入成功，ID: {user_id}")

        # 3. 批量插入
        # batch_data = [
        #     {'name': 'Bob', 'age': 30, 'email': 'bob@example.com'},
        #     {'name': 'Charlie', 'age': 35, 'email': None}
        # ]
        # affected = mysql.batch_insert('users', batch_data)
        # print(f"批量插入影响行数: {affected}")

        # 4.查询数据
        results = mysql.execute_query("SELECT * FROM Student WHERE s_sex = %s and s_id > %s", ("男", 1))
        logger.info(f"查询结果: {results}")

        # 5. 更新数据
        # updated = mysql.execute_update(
        #     "UPDATE users SET email = %s WHERE name = %s",
        #     ('new_email@example.com', 'Alice')
        # )
        # print(f"更新影响行数: {updated}")

        # 6. 事务示例
        # transaction_actions = [
        #     {'sql': "UPDATE users SET age = age + 1 WHERE name = %s", 'params': ('Alice',)},
        #     {'sql': "DELETE FROM users WHERE age > %s", 'params': (40,)}
        # ]
        # success = mysql.execute_transaction(transaction_actions)
        # print(f"事务执行结果: {'成功' if success else '失败'}")

    finally:
        mysql.close()
